from text_processor import TextVectorizer
import json

# 目标电话号码和公司
target_phone = "15937098754"
expected_company = "北京科技创新有限公司"

# 创建并配置向量器
def test_direct_vectorizer():
    # 1. 初始化向量器
    vectorizer = TextVectorizer(use_tfidf=True)
    
    # 2. 准备测试数据
    test_text = """
这是一个测试文件，包含电话号码与公司的关联信息。

以下是一些公司和联系人信息：
1. 北京科技创新有限公司 - 张三 - 15937098754
2. 上海未来信息技术有限公司 - 李四 - 13812345678
3. 广州智慧城市发展有限公司 - 王五 - 13987654321
    """
    
    # 3. 向量化并存储测试数据
    vectorizer.vectorize_and_store(test_text, metadata={"source": "test_phone_company.txt"})
    print("测试数据已加载到向量器")
    
    # 4. 查看词汇表和文档频率
    print("\n=== 词汇表和文档频率 ===")
    print(f"词汇表包含电话号码: {'15937098754' in vectorizer.vocab}")
    print(f"词汇表包含公司名称: {'北京科技创新有限公司' in vectorizer.vocab}")
    print(f"电话号码文档频率: {vectorizer.document_freq.get('15937098754', 0)}")
    print(f"公司名称文档频率: {vectorizer.document_freq.get('北京科技创新有限公司', 0)}")
    
    # 5. 直接查询电话号码
    print("\n=== 查询电话号码结果 ===")
    results = vectorizer.query(target_phone, n_results=3, similarity_metric="cosine")
    print(f"查询结果: {json.dumps(results, ensure_ascii=False, indent=2)}")
    
    # 6. 验证结果
    print("\n=== 验证结果 ===")
    found = False
    for result in results:
        if expected_company in result["text"] and target_phone in result["text"]:
            found = True
            print(f"成功找到关联信息: {expected_company} - {target_phone}")
            print(f"相似度: {result['similarity']:.4f}")
            break
    
    if not found:
        print(f"未找到{target_phone}与{expected_company}的关联信息")
    
    return found

if __name__ == "__main__":
    test_direct_vectorizer()